package com.joker.aigc.ai.config;


import com.joker.aigc.ai.service.ChatAssistant;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;

/**
 * @author : feixiang.li
 * @since : 2025-10-08 14:21
 */
@Configuration
public class LLMConfig {

    /**
     * 向量模型
     */
    @Value("${embedding.base-url}")
    private String aliApiUrl;

    /**
     * 向量模型的apiKey
     */
    @Value("${embedding.api-key}")
    private String aliApiKey;


    @Value("${qdrant.host}")
    private String qdrantHost;
    @Value("${qdrant.port}")
    private int qdrantPort;

    @Value("${qdrant.collection-name}")
    private String qdrantCollectionName;


    @Bean
    public ChatModel chatModel() {
        return OpenAiChatModel.builder()
                .apiKey(aliApiKey)
                .baseUrl(aliApiUrl)
                .modelName("qwen3-max")
                .timeout(Duration.ofSeconds(60))
                .temperature(0.7).logRequests(true)
                .logResponses(true)
                .build();
    }

    @Bean
    public StreamingChatModel streamingChatModel() {
        return OpenAiStreamingChatModel.builder()
                .apiKey(aliApiKey)
                .baseUrl(aliApiUrl)
                .modelName("qwen3-max")
                .timeout(Duration.ofSeconds(60))
                .temperature(0.7).logRequests(true)
                .logResponses(true)
                .build();
    }

    @Bean
    public ChatAssistant assistant(StreamingChatModel streamingChatModel,
                                   EmbeddingStore<TextSegment> embeddingStore,
                                   EmbeddingModel embeddingModel) {
        EmbeddingStoreContentRetriever build = EmbeddingStoreContentRetriever.builder()
                .maxResults(3)
                .minScore(0.7)
                .embeddingStore(embeddingStore)
                .embeddingModel(embeddingModel)
                .build();
        return AiServices.builder(ChatAssistant.class)
                .streamingChatModel(streamingChatModel)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(build)
                .build();
    }

    @Bean
    public EmbeddingModel embeddingModel() {
        return OpenAiEmbeddingModel.builder()
                .apiKey(aliApiKey)
                .baseUrl(aliApiUrl)
                .modelName("text-embedding-v4")
                .timeout(Duration.ofSeconds(60))
                .build();
    }

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return QdrantEmbeddingStore.builder()
                .host(qdrantHost)
                .port(qdrantPort)
                .collectionName(qdrantCollectionName)
                .build();
    }

    @Bean
    public QdrantClient qdrantClient() {
        QdrantGrpcClient.Builder builder =
                QdrantGrpcClient.newBuilder(qdrantHost, qdrantPort, false);
        return new QdrantClient(builder.build());
    }
}
